﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.Net;
using System.Net.NetworkInformation;
using System.Diagnostics;
using System.Collections;

namespace System
{
    [StructLayout(LayoutKind.Sequential)]
    public struct MIB_TCPROW_OWNER_PID
    {
        public TcpState dwState;
        public uint dwLocalAddr;
        public int dwLocalPort;
        public uint dwRemoteAddr;
        public int dwRemotePort;
        public int dwOwningPid;
    }
    [StructLayout(LayoutKind.Sequential)]
    public struct MIB_TCPTABLE_OWNER_PID
    {
        public uint dwNumEntries;
        private MIB_TCPROW_OWNER_PID table;
    }
    [StructLayout(LayoutKind.Sequential)]
    public class MIB_TCPROW
    {
        public int dwState;
        public int dwLocalAddr;
        public int dwLocalPort;
        public int dwRemoteAddr;
        public int dwRemotePort;
    }
    [StructLayout(LayoutKind.Sequential)]
    public struct MIB_TCPTABLE
    {
        public int dwNumEntries;
        public MIB_TCPROW[] table;
    }    

    public class TcpManager
    {
        #region PInvoke define
        private const int TCP_TABLE_OWNER_PID_ALL = 5;

        [DllImport("Iphlpapi.dll")]
        static extern int GetTcpTable(IntPtr pTcpTable, ref int pdwSize, bool bOrder);

        [DllImport("Iphlpapi.dll")]
        static extern int SendARP(Int32 DestIP, Int32 SrcIP, ref Int64 MacAddr, ref Int32 PhyAddrLen);

        [DllImport("Ws2_32.dll")]
        static extern Int32 inet_addr(string ipaddr);

        [DllImport("Ws2_32.dll")]
        static extern ushort ntohs(ushort netshort);

        [DllImport("iphlpapi.dll")]
        static extern int SetTcpEntry(ref MIB_TCPROW pTcpRow);
        [DllImport("iphlpapi.dll")]
        static extern int SetTcpEntry(IntPtr pTcpRow);

        [DllImport("iphlpapi.dll", SetLastError = true)]
        static extern uint GetExtendedTcpTable(
            IntPtr pTcpTable, ref int dwOutBufLen, bool sort, int ipVersion, int tblClass, int reserved);

        #endregion
        public static ArrayList GetAllTcpConnectionInfos()
        {
            ArrayList infos = new ArrayList();

            MIB_TCPTABLE tcpTableData = new MIB_TCPTABLE();
            tcpTableData = Native.GetTcpTableInfo();

            foreach (MIB_TCPROW row in tcpTableData.table)
            {
                try
                {
                    int localPort = TranslatePort(row.dwLocalPort);
                    int remotePort = TranslatePort(row.dwRemotePort);

                    TcpConnectionInfo conn =
                        new TcpConnectionInfo(

                            new IPEndPoint(row.dwLocalAddr, localPort),

                            new IPEndPoint(row.dwRemoteAddr, remotePort),

                           (TcpState)row.dwState);

                    infos.Add(conn);
                }
                catch
                { }
            }

            return infos;
        }

        public static TcpConnectionInfo[] GetTableByProcess(int pid)
        {
            MIB_TCPROW_OWNER_PID[] tcpRows = GetAllTcpConnections();
            if (tcpRows == null) return null;
            List<TcpConnectionInfo> list = new List<TcpConnectionInfo>();
            foreach (MIB_TCPROW_OWNER_PID row in tcpRows)
            {
                if (row.dwOwningPid == pid)
                {
                    int localPort = TranslatePort(row.dwLocalPort);
                    int remotePort = TranslatePort(row.dwRemotePort);
                    TcpConnectionInfo conn =
                        new TcpConnectionInfo(
                            new IPEndPoint(row.dwLocalAddr, localPort),
                            new IPEndPoint(row.dwRemoteAddr, remotePort),
                            row.dwState);
                    list.Add(conn);
                }
            }
            return list.ToArray();
        }

        private static IntPtr Get(object obj)
        {
            IntPtr ptr = Marshal.AllocCoTaskMem(Marshal.SizeOf(obj));
            Marshal.StructureToPtr(obj, ptr, false);
            return ptr;
        }
        //public static int KillByPort(int port)
        ////{
        ////    MIB_TCPROW[] rows = Native.GetTcpTableInfo().table;
        ////    int i = 0;
        ////    foreach (MIB_TCPROW row in rows)
        ////    {
        ////       // if (row.dwLocalPort == port)
        ////       // {
        ////            row.dwState = (int)TcpState.DeleteTcb;
        ////            var r = row;
        ////              IntPtr pr = Get(row);
        ////            int sr = SetTcpEntry(pr);
        ////            if (sr == 0)
        ////                i++;
        ////       // }
        ////    }
        ////    return i;
        //}
        public static bool Kill(MIB_TCPROW row)
        {
            row.dwState = (int)TcpState.DeleteTcb;
            IntPtr pr = Get(row);
            int s = SetTcpEntry(pr);
            ZLog.SaveLog(Native.GetIpAddress(row.dwLocalAddr) + ":" + row.dwLocalPort.ToString(), Native.GetIpAddress(row.dwRemoteAddr) + ":" + row.dwRemotePort.ToString(), s.ToString());
            return s == 0;
        }

        public static bool Kill(TcpConnectionInfo conn)
        {
            if (conn == null)
                throw new ArgumentNullException("conn");
            MIB_TCPROW row = new MIB_TCPROW();
            row.dwState = Convert.ToInt32(TcpState.DeleteTcb);
#pragma warning disable 612,618
            row.dwLocalAddr = (int)conn.LocalEndPoint.Address.Address;
#pragma warning restore 612,618
            row.dwLocalPort = TranslatePort(conn.LocalEndPoint.Port);
#pragma warning disable 612,618
            row.dwRemoteAddr = (int)conn.RemoteEndPoint.Address.Address;
#pragma warning restore 612,618
            row.dwRemotePort = TranslatePort(conn.RemoteEndPoint.Port);
            int s = SetTcpEntry(ref row);
            ZLog.SaveLog(conn.LocalEndPoint.Address.ToString() + ":" + conn.LocalEndPoint.Port.ToString(), conn.RemoteEndPoint.Address.ToString() + ":" + conn.RemoteEndPoint.Port.ToString(), s.ToString());
            return s == 0;
        }

        public static bool Kill(IPEndPoint localEndPoint, IPEndPoint remoteEndPoint)
        {
            if (localEndPoint == null) throw new ArgumentNullException("localEndPoint");
            if (remoteEndPoint == null) throw new ArgumentNullException("remoteEndPoint");
            MIB_TCPROW row = new MIB_TCPROW();
            row.dwState = Convert.ToInt32(TcpState.DeleteTcb);
#pragma warning disable 612,618
            row.dwLocalAddr = (int)localEndPoint.Address.Address;
#pragma warning restore 612,618
            row.dwLocalPort = TranslatePort(localEndPoint.Port);
#pragma warning disable 612,618
            row.dwRemoteAddr = (int)remoteEndPoint.Address.Address;
#pragma warning restore 612,618
            row.dwRemotePort = TranslatePort(remoteEndPoint.Port);
            int s = SetTcpEntry(ref row);
            ZLog.SaveLog(localEndPoint.Address.ToString() + ":" + localEndPoint.Port.ToString(), remoteEndPoint.Address.ToString() + ":" + remoteEndPoint.Port.ToString(), s.ToString());
            return s == 0;
        }

        public static TcpConnectionInfo[] GetTalbeByCurrentProcess()
        {
            return GetTableByProcess(Process.GetCurrentProcess().Id);
        }
        private static MIB_TCPROW_OWNER_PID[] GetAllTcpConnections()
        {
            const int NO_ERROR = 0;
            const int IP_v4 = 2;
            MIB_TCPROW_OWNER_PID[] tTable = null;
            int buffSize = 0;
            GetExtendedTcpTable(IntPtr.Zero, ref buffSize, true, IP_v4, TCP_TABLE_OWNER_PID_ALL, 0);
            IntPtr buffTable = Marshal.AllocHGlobal(buffSize);
            try
            {
                if (NO_ERROR != GetExtendedTcpTable(buffTable, ref buffSize, true, IP_v4, TCP_TABLE_OWNER_PID_ALL, 0)) return null;
                MIB_TCPTABLE_OWNER_PID tab =
                    (MIB_TCPTABLE_OWNER_PID)Marshal.PtrToStructure(buffTable, typeof(MIB_TCPTABLE_OWNER_PID));
                IntPtr rowPtr = (IntPtr)((long)buffTable + Marshal.SizeOf(tab.dwNumEntries));
                tTable = new MIB_TCPROW_OWNER_PID[tab.dwNumEntries];

                int rowSize = Marshal.SizeOf(typeof(MIB_TCPROW_OWNER_PID));
                for (int i = 0; i < tab.dwNumEntries; i++)
                {
                    MIB_TCPROW_OWNER_PID tcpRow =
                        (MIB_TCPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_TCPROW_OWNER_PID));
                    tTable[i] = tcpRow;
                    rowPtr = (IntPtr)((int)rowPtr + rowSize);
                }
            }
            finally
            {
                Marshal.FreeHGlobal(buffTable);
            }
            return tTable;
        }

        private static int TranslatePort(int port)
        {
            return ((port & 0xFF) << 8 | (port & 0xFF00) >> 8);
        }
    }
}